package link.luyu.plugin.brochain.crypto;

import static org.web3j.utils.Assertions.verifyPrecondition;

import java.math.BigInteger;
import java.util.Arrays;
import org.bouncycastle.asn1.x9.X9ECParameters;
import org.bouncycastle.asn1.x9.X9IntegerConverter;
import org.bouncycastle.crypto.ec.CustomNamedCurves;
import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.math.ec.ECAlgorithms;
import org.bouncycastle.math.ec.ECConstants;
import org.bouncycastle.math.ec.ECPoint;
import org.bouncycastle.math.ec.FixedPointCombMultiplier;

/** The SM2 Digital Signature algorithm logic. */
public class SM2 {

    public static final X9ECParameters CURVE_PARAMS = CustomNamedCurves.getByName("sm2p256v1");
    static final ECDomainParameters CURVE =
            new ECDomainParameters(
                    CURVE_PARAMS.getCurve(),
                    CURVE_PARAMS.getG(),
                    CURVE_PARAMS.getN(),
                    CURVE_PARAMS.getH()); // TODO Need H?
    static final BigInteger HALF_CURVE_ORDER = CURVE_PARAMS.getN().shiftRight(1);

    public static SM2Signature signMessage(byte[] message, SM2KeyPair keyPair) {
        return signMessage(message, keyPair, true);
    }

    public static SM2Signature signMessage(byte[] message, SM2KeyPair keyPair, boolean needToHash) {
        byte[] messageHash;
        if (needToHash) {
            messageHash = Hash.sm3(message);
        } else {
            messageHash = message;
        }

        return keyPair.sign(messageHash);
    }

    public static BigInteger recoverFromSignature(SM2Signature signature, byte[] message) {
        return recoverFromSignature(
                signature.getV(), signature.getR(), signature.getS(), message, true);
    }

    public static BigInteger recoverFromSignature(
            SM2Signature signature, byte[] message, boolean needToHash) {
        return recoverFromSignature(
                signature.getV(), signature.getR(), signature.getS(), message, needToHash);
    }

    public static BigInteger recoverFromSignature(
            byte v, BigInteger r, BigInteger s, byte[] message, boolean needToHash) {
        verifyPrecondition(v >= 0, "v must be positive");
        verifyPrecondition(r.signum() >= 0, "r must be positive");
        verifyPrecondition(s.signum() >= 0, "s must be positive");
        verifyPrecondition(message != null, "message cannot be null");

        byte[] messageHash;
        if (needToHash) {
            messageHash = Hash.sm3(message);
        } else {
            messageHash = message;
        }

        BigInteger n = CURVE.getN();

        if (r.compareTo(ECConstants.ONE) < 0 || r.compareTo(n) >= 0) {
            return null;
        }

        if (s.compareTo(ECConstants.ONE) < 0 || s.compareTo(n) >= 0) {
            return null;
        }

        BigInteger e = new BigInteger(1, messageHash);

        BigInteger t = r.add(s).mod(n);
        if (t.equals(ECConstants.ZERO)) {
            return null;
        }

        BigInteger rx = r.subtract(e).mod(n);
        ECPoint pt = decompressKey(rx, (v & 1) == 1);
        BigInteger rr = t.modInverse(n);
        BigInteger g = s.multiply(rr).mod(n);
        g = g.negate().mod(n);
        ECPoint q = ECAlgorithms.sumOfTwoMultiplies(CURVE.getG(), g, pt, rr);

        byte[] qBytes = q.getEncoded(false);
        return new BigInteger(1, Arrays.copyOfRange(qBytes, 1, qBytes.length));
    }

    /** Decompress a compressed public key (x co-ord and low-bit of y-coord). */
    private static ECPoint decompressKey(BigInteger xBN, boolean yBit) {
        X9IntegerConverter x9 = new X9IntegerConverter();
        byte[] compEnc = x9.integerToBytes(xBN, 1 + x9.getByteLength(CURVE.getCurve()));
        compEnc[0] = (byte) (yBit ? 0x03 : 0x02);
        return CURVE.getCurve().decodePoint(compEnc);
    }

    /**
     * Returns public key from the given private key.
     *
     * @param privateKey the private key to derive the public key from
     * @return BigInteger encoded public key
     */
    public static BigInteger publicKeyFromPrivate(BigInteger privateKey) {
        ECPoint point = publicPointFromPrivate(privateKey);

        byte[] encoded = point.getEncoded(false);
        return new BigInteger(1, Arrays.copyOfRange(encoded, 1, encoded.length)); // remove prefix
    }

    /**
     * Returns public key point from the given private key.
     *
     * @param privateKey the private key to derive the public key from
     * @return ECPoint public key
     */
    public static ECPoint publicPointFromPrivate(BigInteger privateKey) {
        /*
         * TODO: FixedPointCombMultiplier currently doesn't support scalars longer than the group
         * order, but that could change in future versions.
         */
        if (privateKey.bitLength() > CURVE.getN().bitLength()) {
            privateKey = privateKey.mod(CURVE.getN());
        }
        return new FixedPointCombMultiplier().multiply(CURVE.getG(), privateKey);
    }
}
